`include "defines.v"

module MEM_Stage (
    input  wire [63: 0] pc_addr_i,
    output wire [63: 0] pc_addr_o,

    input  wire [`xlen] ALU_out_i,
    output wire [`xlen] ALU_out_o,

    input  wire [ 4: 0] rd_addr_i,
    output wire [ 4: 0] rd_addr_o,

    input  wire [11: 0] csr_addr_i,
    output wire [11: 0] csr_addr_o,

    input  wire [`xlen] data_to_mem_i,
    input  wire [`xlen] data_from_mem_i,
    output wire [`xlen] data_to_mem_o,
    output wire [`xlen] data_from_mem_o,

    input  wire         jump,
    input  wire [ 2: 0] branch,
    input  wire [ 3: 0] flags,
    output wire         pcsel,

    input  wire         is_system_inst_i,
    input  wire [ 3: 0] MemWr_i,
    input  wire         RegWr_i,
    input  wire         MemToReg_i,
    input  wire [ 2: 0] load_ext_i,
    output wire         is_system_inst_o,
    output wire         MemWr_ena,
    output reg  [ 7: 0] mem_byte_ena,
    output wire         RegWr_o,
    output wire         MemToReg_o,
    output wire         MemRd,
    output wire [ 2: 0] load_ext_o
);

assign pc_addr_o = pc_addr_i;
assign ALU_out_o = is_system_inst_i ? data_to_mem_i : ALU_out_i;
assign rd_addr_o = rd_addr_i;

assign data_to_mem_o = data_to_mem_i;
assign data_from_mem_o = is_system_inst_i ? ALU_out_i : data_from_mem_i;

assign RegWr_o = RegWr_i;
assign MemToReg_o = MemToReg_i;
assign MemRd = MemToReg_o;
assign load_ext_o = load_ext_i;
assign is_system_inst_o = is_system_inst_i;

//byte ena generate
assign MemWr_ena = MemWr_i[2];
always @(*) begin
    if (MemWr_ena) begin
        case (MemWr_i[1:0])
        2'b00 : mem_byte_ena = 8'b0000_0001;
        2'b01 : mem_byte_ena = 8'b0000_0011;
        2'b10 : mem_byte_ena = 8'b0000_1111;
        2'b11 : mem_byte_ena = 8'b1111_1111;
        default: mem_byte_ena = 8'b0000_0001;
        endcase
    end
    else begin
        case (load_ext_o)
        `SEXT8  : mem_byte_ena = 8'b0000_0001;
        `SEXT16 : mem_byte_ena = 8'b0000_0011;
        `SEXT32 : mem_byte_ena = 8'b0000_1111;
        `ZEXT8  : mem_byte_ena = 8'b0000_0001;
        `ZEXT16 : mem_byte_ena = 8'b0000_0011;
        `ZEXT32 : mem_byte_ena = 8'b0000_1111;
        `NOEXT  : mem_byte_ena = 8'b1111_1111;
        default : mem_byte_ena = 8'b0000_0001;
        endcase
    end
end
//

//pcsel
reg branch_result;

always @(*) begin
    case (branch)
        `notbranch : begin
            branch_result = 1'b0;
        end
        `eq : begin
            branch_result = `ZF;
        end
        `ne : begin
            branch_result = ~`ZF;
        end
        `lt : begin
            branch_result = `SF^`OF;
        end
        `ge : begin
            branch_result = `SF^~`OF | `ZF;
        end
        `ltu : begin
            branch_result = `CF;
        end
        `geu : begin
            branch_result = ~`CF | `ZF;
        end
        default : branch_result = 1'b0;
    endcase
end

assign pcsel = jump | branch_result;
//



endmodule